import os
from torchvision import transforms
import torch, PIL
import torch
from PIL import Image

means = [0.485, 0.456, 0.406]
stds = [0.229, 0.224, 0.225]
img_transforms = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(means,stds)])

t_stds = torch.tensor(stds).cuda().half()[:,None,None]
t_means = torch.tensor(means).cuda().half()[:,None,None]

def tensor2im(var):
    return var.mul(t_stds).add(t_means).mul(255.).clamp(0, 255).permute(1, 2, 0)

def proc_pil_img(input_image, model):
    transformed_image = img_transforms(input_image)[None, ...].cuda().half()

    with torch.no_grad():
        result_image = model(transformed_image)[0]
        print(result_image.shape)
        output_image = tensor2im(result_image)
        output_image = output_image.detach().cpu().numpy().astype('uint8')
        output_image = PIL.Image.fromarray(output_image)
    return output_image

def file_name(file_dir):
    img_path_list = []
    for root, dirs, files in os.walk(file_dir):

        for file in files:
             img_path_list.append((os.path.join(root, file),file))
    return img_path_list

if __name__ == '__main__':
    # 手动指定图像路径
    img_path = './testJPG/3073.png'  # 指定你要处理的图像路径
    save_dir = './resultJPG'  # 结果保存目录
    os.makedirs(save_dir, exist_ok=True)  # 创建保存目录

    # 加载模型
    modelv4 = torch.jit.load('./models/ArcaneGANv0.4.jit').eval().cuda().half()

    # 处理图像
    im = Image.open(img_path).convert('RGB')  # 打开指定的图像
    save_path = os.path.join(save_dir, 'processed_sample.jpg')  # 结果保存路径，指定文件名

    # 处理并保存图像
    res = proc_pil_img(im, modelv4)
    res.save(save_path)  # 保存处理后的图像

    print(f"处理后的图像已保存到: {save_path}")



